import pandas as pd
import talib
import tushare as ts
import time


def get_realtime_stock_data(stock_code):
    """
    获取实时股票数据
    :param stock_code: 股票代码
    :return: 股票实时行情DataFrame
    """
    return ts.get_realtime_quotes(stock_code)


def calculate_tangguangtrend(df):
    """
    计算唐奇安趋势
    :param df: 股票数据DataFrame
    :return: 计算结果DataFrame
    """
    # 计算唐奇安上限和下限
    df['upper'], df['middle'], df['lower'] = talib.BBANDS(df['close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)

    # 计算唐奇安趋势
    df['sma'] = talib.SMA(df['close'], timeperiod=20)
    df['std'] = talib.STDDEV(df['close'], timeperiod=20)
    df['upper'] = df['sma'] + 2 * df['std']
    df['lower'] = df['sma'] - 2 * df['std']

    # 返回计算结果
    return df


# 读取股票数据
df = pd.read_csv('stock_data.csv')

while True:
    # 获取上证指数和某只股票的实时行情数据
    index_realtime_data = get_realtime_stock_data('sh')
    stock_realtime_data = get_realtime_stock_data('000001')

    # 计算唐奇安趋势
    df = calculate_tangguangtrend(df)

    # 输出计算结果和实时行情
    print("实时行情数据：")
    print(stock_realtime_data)
    print("唐奇安趋势计算结果：")
    print(df.tail())

    # 等待三十秒再继续获取数据
    time.sleep(30)